
# 脚本功能：创建适合于deepseek的集合；实际执行结果是 可以连接并成功创建适合于deepseek要求的集合
# 读取本地文件进入集合

from config import envConfig, config
from langchain_ollama import OllamaEmbeddings
import logging
from pymilvus import (
    connections,
    utility,
    Collection,
    FieldSchema,
    CollectionSchema,
    DataType,
    MilvusException
)
import time

logger = logging.getLogger('rebuild_collection')


def create_collection_directly(collection_name, dim=1536):
    """直接使用pymilvus创建集合"""
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name="metadata", dtype=DataType.JSON)
    ]

    # 创建schema
    schema = CollectionSchema(fields, description="使用deepseek-r1:1.5b嵌入模型")

    # 创建集合
    collection = Collection(
        name=collection_name,
        schema=schema,
        using="default"
    )

    # 创建索引
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}
    }
    collection.create_index("embedding", index_params)
    return collection


def rebuild_milvus_collection():
    """使用原生pymilvus API重建集合"""
    # 配置连接
    milvus_conn = {
        "host": envConfig.MILVUS_LOCAL_URI.split("://")[-1].split(":")[0],
        "port": "19530",
        "user": envConfig.MILVUS_USER,
        "password": envConfig.MILVUS_PASSWORD
    }

    # 测试嵌入模型
    embeddings = OllamaEmbeddings(
        model="deepseek-r1:1.5b",
        base_url=envConfig.OLLAMA_URL
    )
    try:
        test_embed = embeddings.embed_query("test")
        dim = len(test_embed)
        logger.info(f"嵌入模型维度: {dim}")
    except Exception as e:
        logger.error(f"嵌入模型测试失败: {e}")
        return False

    collection_name = config.APP_NAME.replace(' ', '_')

    try:
        # 连接Milvus
        connections.connect(
            alias="default",
            host=milvus_conn["host"],
            port=milvus_conn["port"],
            user=milvus_conn["user"],
            password=milvus_conn["password"],
            secure=False
        )

        # 删除现有集合
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
            time.sleep(1)

        # 直接创建集合
        collection = create_collection_directly(collection_name, dim)
        logger.info(f"集合 {collection_name} 创建成功")

        # 验证
        if not utility.has_collection(collection_name):
            raise MilvusException("集合验证失败")

        return True

    except Exception as e:
        logger.error(f"重建失败: {e}")
        return False
    finally:
        connections.disconnect("default")


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    print("=== Milvus集合重建程序 ===")
    start_time = time.time()

    if rebuild_milvus_collection():
        print(f"\n重建成功! 耗时: {time.time() - start_time:.2f}s")
    else:
        print("\n重建失败，请检查日志")

    print("=" * 30)
